/*
 * Copyright (c) Extreme Vision Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef YOLOv7_PREPOSTPROCESSOR_H
#define YOLOv7_PREPOSTPROCESSOR_H

#include "opencv2/core.hpp"
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

#include "evdeploy/core/ev_common_data.h"
#include "evdeploy/cv/common/common.h"

namespace ev
{
    namespace vision
    {

        class YOLOv7Preprocessor
        {
        public:
            YOLOv7Preprocessor();

            ~YOLOv7Preprocessor();
            /**
             * @brief preprocessing for YOLOv7
             * @param[in] in_mat: input cv::Mat
             * @param[in] out_width: resized Mat Size(out_width, out_width)
             * @param[in] normalize: true for normalize to 1, false for without normalization, default true
             * @param[in] bgr2rgb: true for bgr2rgb, false for bgr format, default true
             * @param[out] out_mat: output cv::Mat
             * @return EVStatus
             */
            EVStatus Run(cv::Mat &in_mat, cv::Mat &out_mat, int out_width, bool normalize = true, bool bgr2rgb = true);
            /**
             * @brief get scale from preprocessing, used for postprocessing, e.g., compute the coordinates in original image
             * @return scale
             */
            float GetScale();

        private:
            float m_scale;
        };

        class YOLOv7Postprocessor
        {
        public:
            YOLOv7Postprocessor();

            ~YOLOv7Postprocessor();
            /**
             * @brief postprocessing for YOLOv7
             * @param[in] out: output data from runtime inference
             * @param[in] scale: scale from preprocessing
             * @param[in] thresh: thresh for detection confidence
             * @param[in] img_w: original image width
             * @param[in] img_h: original image height
             * @param[in] nms_thresh: thresh for nms
             * @param[in] class_aware whether to distinguish labels in nms
             * @param[out] objects: detection output wrapped in BoxInfo array
             * @return EVStatus
             */
            EVStatus Run(EVMatData *out, std::vector<ev::vision::BoxInfo> &objects, float scale, float thresh, const int img_w, const int img_h, float nms_thresh = 0.45, bool class_aware = false);
        };

    } // namespace vision
} // namespace ev

#endif
